import pandas as pd
from transformer.transformer import create_model_1
import tensorflow as tf
import numpy as np
from tensorflow import keras

#tf.config.run_functions_eagerly("True")

df = pd.read_csv("tongji/dynamic_normalization/1-bt.csv")
train = pd.read_csv("tongji/static/train_case.csv")

df = pd.merge(df, train, how="inner", on=["caseid"])

def generate_dynamic(db, ob_win, features):
    #choose the features
    dynamic_data = db[features].values
    #get the length of row
    data_length = dynamic_data.shape[0]
    #return the data
    for start,stop in zip(range(0, data_length-ob_win), range(ob_win, data_length)):
        yield dynamic_data[start:stop, :]

grouped = df.groupby("caseid")
print("...")
dynamic = [list(generate_dynamic(group, 30, ["MAP", "SDP", "HR", "stage"])) for value, group in grouped]
#drop []
dynamic = [dd for dd in dynamic if dd != []]
#type list to type array
dynamic = np.concatenate(dynamic).astype(np.float32)
print(dynamic[0])

dynamic = dynamic[0:512, :, :]

print("....")
dynamic_input = dynamic[:, 0:15, 0:3]
print("intput: ", dynamic_input.shape)

dynamic_output = dynamic[:, 15:29, 3:]
d_z = np.zeros(shape=(dynamic_output.shape[0], 1, 1))
dynamic_output = np.concatenate((d_z, dynamic_output), axis=1)
print("output: ", dynamic_output.shape)

label = dynamic[:, 15:30, 3:]
label = np.squeeze(label)
#label = tf.one_hot(label, depth=5)
print("label: ", label.shape)


model = create_model_1(dynamic_input.shape[1:])
model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['acc'])
model.summary()

history = model.fit([dynamic_input, dynamic_output], label, epochs=200, batch_size=256, validation_split=0.3, verbose=1, 
                    callbacks=[keras.callbacks.EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=0, mode='min'),
                               keras.callbacks.ModelCheckpoint("trans_2.h5", monitor='val_loss', save_best_only=True, mode='min', verbose=0)])
